
import numpy as np
import matplotlib.pyplot as plt
from ruamel.yaml import YAML
from BeaverFramework.Framework.model import Model
from BeaverFramework.Framework.assembly import Assembly
from BeaverFramework.Framework.framework import Framework
from BeaverFramework.Framework.utles import load_config
from BeaverFramework.Models.SpringDamper.SpringDamper import MassSpringDamperModel
from BeaverFramework.Models.SpringDamper.SpringDamper import PostProcessingModel
from os.path import join
import os
from dotenv import load_dotenv

"""
BeaverFramework Spring-Damper Example Module

This module demonstrates the usage of BeaverFramework by implementing a mass-spring-damper
system simulation. It showcases framework features including model definition, assembly,
and execution.

Authors:
    MichaelZhang

Key Features:
    - Complete mass-spring-damper system simulation
    - Framework initialization and configuration
    - Model assembly and linking
    - Post-processing and visualization

Required Setup:
    - .env file with RUN_FRAME_BASE configuration
    - Log directory for output logging
    - YAML configuration file for simulation parameters

Usage:
    python SpringDamper.py
"""

def run_simulation():
    """
    Execute a complete mass-spring-damper system simulation.

    Required varables and files:
        :param - .env file with RUN_FRAME_BASE=/home/mich/apps/AIcoding/beaver-framework
        :param - Log file for output logging
        :param - config/msd_simulation_input.yaml for input variables

    Required Environment:
        - RUN_FRAME_BASE environment variable pointing to framework base directory
        - Log directory for output files
        - Configuration YAML file with simulation parameters

    Configuration Parameters (from YAML):
        mass_spring_damper:
          mass: 1.0  # kg
          spring_constant: 10.0  # N/m
          damping_coefficient: 0.5  # Ns/m
          initial_position: 1.0  # m
          initial_velocity: 0.0  # m/s
          simulation_time: 10.0  # s
          time_step: 0.01  # s

    Raises:
        FileNotFoundError: If configuration or log files cannot be accessed
        EnvironmentError: If required environment variables are not set
    """
    # Load environment variables and set paths
    load_dotenv()
    run_base = join(os.getenv("RUN_FRAME_BASE"),'Applications')
    log_file = join(os.getenv("RUN_FRAME_BASE"),'Log/SpringDampler.log')
    if not os.path.exists(log_file):
        os.makedirs(log_file)
    # Load configuration
    config = load_config(join(run_base,'config/msd_simulation_input.yaml'))

    # Initialize simulation models
    msd_model = MassSpringDamperModel("MSD", config['mass_spring_damper'])
    post_processing_model = PostProcessingModel("PostProcessing", config['post_processing'])

    # Create and configure assembly
    assembly = Assembly("MSDSimulation")
    assembly.add_model(msd_model)
    assembly.add_model(post_processing_model)

    # Link model inputs and outputs
    post_processing_model.link_input('time', msd_model, 'time')
    post_processing_model.link_input('results', msd_model, 'results')
    msd_model.link_output('time', post_processing_model, 'time')
    msd_model.link_output('results', post_processing_model, 'results')

    # Initialize framework and execute simulation
    framework = Framework()
    framework.setup(assembly)

    # Run the framework
    framework.run()

    # Show the plot
    plt.show()

if __name__ == "__main__":
    """
    Main entry point for the spring-damper example simulation.
    
    Executes the simulation and handles any top-level exceptions.
    """
    try:
        run_simulation()
    except Exception as e:
        print(f"Simulation failed: {str(e)}")
        raise
